"""
手指检测模块，使用ultralytics YOLO模型检测和识别手指位置。
"""
import cv2
import numpy as np
import os
import math
from typing import List, Tuple, Dict, Optional, Union
from dataclasses import dataclass
from ultralytics import YOLO


@dataclass
class FingerDetection:
    """手指检测结果数据类"""
    # 手指关键点坐标 (x, y) - 归一化坐标 (0-1)
    landmarks: List[Tuple[float, float]]
    # 手指类型 (0: 拇指, 1: 食指, 2: 中指, 3: 无名指, 4: 小指)
    finger_type: int
    # 置信度
    confidence: float
    # 指尖坐标 (x, y) - 像素坐标
    tip_position: Tuple[int, int]
    # 延伸方向向量 (dx, dy) - 标准化的方向向量
    direction: Tuple[float, float]
    # 是否为左手
    is_left_hand: bool


class FingerDetector:
    """
    手指检测器类，使用ultralytics YOLO模型检测和跟踪手指。
    优化了左右手的处理和指向计算。
    """
    
    # 手指类型映射
    FINGER_TYPES = {
        'thumb': 0,     # 拇指
        'index': 1,     # 食指
        'middle': 2,    # 中指
        'ring': 3,      # 无名指
        'pinky': 4      # 小指
    }
    
    def __init__(
        self,
        model_path: Optional[str] = None,
        confidence_threshold: float = 0.5,
        device: Optional[str] = None
    ):
        """
        初始化手指检测器
        
        参数:
            model_path: YOLO模型路径，如果为None则使用默认路径
            confidence_threshold: 置信度阈值
            device: 运行设备，可以是'cuda'或'cpu'，如果为None则自动选择
        """
        self.confidence_threshold = confidence_threshold
        
        # 设置运行设备
        if device is None:
            import torch
            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        else:
            self.device = device
            
        # 加载模型
        if model_path is None:
            # 如果未指定模型路径，则查找默认路径
            default_paths = [
                os.path.join(os.path.dirname(__file__), '../models/yolov11_hand.pt'),
                os.path.join(os.path.dirname(__file__), '../../models/yolov11_hand.pt'),
                'yolov11_hand.pt'
            ]
            
            for path in default_paths:
                if os.path.exists(path):
                    model_path = path
                    break
        
        try:
            # 尝试加载指定的模型
            if model_path is not None and os.path.exists(model_path):
                self.model = YOLO(model_path)
                print(f"成功加载手部YOLO模型 {model_path}，运行设备: {self.device}")
            else:
                # 如果没有找到指定模型，使用预训练模型
                print("未找到指定手部YOLO模型，加载预训练模型...")
                self.model = YOLO("yolov8n-pose.pt")  # 加载小型预训练姿态检测模型
                print(f"已加载预训练YOLO姿态检测模型，运行设备: {self.device}")
            
            # 配置模型运行设备
            self.model = self.model.to(self.device)
            
            # 获取类别列表
            self.class_names = self.model.names if hasattr(self.model, 'names') else {}
            print(f"检测类别: {self.class_names}")
            
        except Exception as e:
            # 如果加载失败，使用备用方法
            print(f"加载YOLO手部检测模型失败: {str(e)}，将使用备用检测方法")
            self.model = None

    def detect_fingers(self, image: np.ndarray) -> List[FingerDetection]:
        """
        检测图像中的手指
        
        参数:
            image: 输入图像
            
        返回:
            List[FingerDetection]: 检测到的手指列表
        """
        if image is None or image.size == 0:
            return []
            
        # 如果模型为None，使用备用方法
        if self.model is None:
            return self._detect_fingers_fallback(image)
            
        # 使用YOLO进行检测
        try:
            # 姿态检测模式
            results = self.model(image, conf=self.confidence_threshold, verbose=False)
            h, w = image.shape[:2]
            
            finger_detections = []
            
            # 处理检测结果
            for result in results:
                # 检查是否有姿态关键点
                if hasattr(result, 'keypoints') and result.keypoints is not None:
                    # 处理姿态关键点
                    keypoints = result.keypoints
                    
                    # 尝试从关键点中检测手指
                    # YOLO姿态模型通常有17个关键点，其中包括手腕
                    # 我们可以使用手腕和手肘的位置来估计手的方向
                    for i in range(len(keypoints)):
                        # 获取手腕关键点 (通常索引为9右腕和10左腕)
                        if len(keypoints) > 0 and len(keypoints[i]) > 0:
                            try:
                                # 获取右手腕和左手腕关键点
                                right_wrist_idx, left_wrist_idx = 9, 10
                                
                                # 检测到右手腕
                                if right_wrist_idx < len(keypoints[i]) and keypoints[i][right_wrist_idx].conf > self.confidence_threshold:
                                    right_wrist = keypoints[i][right_wrist_idx]
                                    right_elbow_idx = 7  # 右肘
                                    
                                    if right_elbow_idx < len(keypoints[i]) and keypoints[i][right_elbow_idx].conf > self.confidence_threshold:
                                        right_elbow = keypoints[i][right_elbow_idx]
                                        
                                        # 构建右手指检测
                                        finger_detection = self._create_finger_detection(
                                            right_wrist, right_elbow, w, h, is_left_hand=False
                                        )
                                        finger_detections.append(finger_detection)
                                
                                # 检测到左手腕
                                if left_wrist_idx < len(keypoints[i]) and keypoints[i][left_wrist_idx].conf > self.confidence_threshold:
                                    left_wrist = keypoints[i][left_wrist_idx]
                                    left_elbow_idx = 8  # 左肘
                                    
                                    if left_elbow_idx < len(keypoints[i]) and keypoints[i][left_elbow_idx].conf > self.confidence_threshold:
                                        left_elbow = keypoints[i][left_elbow_idx]
                                        
                                        # 构建左手指检测
                                        finger_detection = self._create_finger_detection(
                                            left_wrist, left_elbow, w, h, is_left_hand=True
                                        )
                                        finger_detections.append(finger_detection)
                            except Exception as e:
                                print(f"处理手部关键点时出错: {str(e)}")
                                continue
                else:
                    # 如果没有姿态关键点，尝试检测手的边界框
                    boxes = result.boxes
                    for i, box in enumerate(boxes):
                        try:
                            # 获取边界框
                            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
                            w_box, h_box = x2 - x1, y2 - y1
                            
                            # 获取置信度
                            conf = float(box.conf[0].cpu().numpy())
                            
                            # 获取类别ID和名称
                            cls = int(box.cls[0].cpu().numpy())
                            cls_name = self.class_names.get(cls, "")
                            
                            # 检查是否为手部类别
                            if cls_name.lower() in ['hand', 'hands', '手']:
                                # 估计手的中心点和方向
                                center_x, center_y = x1 + w_box // 2, y1 + h_box // 2
                                
                                # 假设手指指向上方
                                dx, dy = 0, -1
                                
                                # 创建手指检测结果
                                # 根据边界框位置猜测是左手还是右手
                                is_left = center_x < (w // 2)
                                
                                landmarks = [
                                    (center_x / w, center_y / h),  # 手心
                                    ((center_x + dx * 20) / w, (center_y + dy * 20) / h)  # 假设的指尖
                                ]
                                
                                finger_detection = FingerDetection(
                                    landmarks=landmarks,
                                    finger_type=1,  # 假设为食指
                                    confidence=conf,
                                    tip_position=(center_x, center_y - h_box // 3),  # 假设指尖在手部上方
                                    direction=(dx, dy),
                                    is_left_hand=is_left
                                )
                                
                                finger_detections.append(finger_detection)
                        except Exception as e:
                            print(f"处理手部边界框时出错: {str(e)}")
                            continue
            
            # 如果没有检测到手指，使用备用方法
            if not finger_detections:
                return self._detect_fingers_fallback(image)
                
            return finger_detections
        except Exception as e:
            print(f"YOLO手部检测出错: {str(e)}")
            return self._detect_fingers_fallback(image)

    def _create_finger_detection(self, wrist, elbow, img_width, img_height, is_left_hand=False):
        """
        根据手腕和手肘关键点创建手指检测结果
        """
        # 获取关键点坐标
        wrist_x, wrist_y = wrist.xy[0].cpu().numpy()
        elbow_x, elbow_y = elbow.xy[0].cpu().numpy()
        
        # 计算手腕到手肘的方向向量
        dx = wrist_x - elbow_x
        dy = wrist_y - elbow_y
        
        # 标准化方向向量
        magnitude = math.sqrt(dx**2 + dy**2)
        if magnitude > 0:
            dx, dy = dx / magnitude, dy / magnitude
            
            # 对于左右手，调整指向方向
            if is_left_hand:
                # 左手，稍微向左旋转
                dx, dy = dx * 0.8 - dy * 0.6, dx * 0.6 + dy * 0.8
            else:
                # 右手，稍微向右旋转
                dx, dy = dx * 0.8 + dy * 0.6, -dx * 0.6 + dy * 0.8
                
        else:
            # 如果无法计算方向，假设指向上方
            dx, dy = 0, -1
            
        # 计算指尖坐标（从手腕沿方向延伸）
        tip_x = int(wrist_x + dx * 30)  # 延伸30个像素
        tip_y = int(wrist_y + dy * 30)
        
        # 确保坐标在图像范围内
        tip_x = max(0, min(tip_x, img_width - 1))
        tip_y = max(0, min(tip_y, img_height - 1))
        
        # 创建检测结果
        landmarks = [
            (wrist_x / img_width, wrist_y / img_height),
            (elbow_x / img_width, elbow_y / img_height),
            (tip_x / img_width, tip_y / img_height)
        ]
        
        return FingerDetection(
            landmarks=landmarks,
            finger_type=1,  # 假设为食指
            confidence=float(wrist.conf.cpu().numpy()),
            tip_position=(tip_x, tip_y),
            direction=(dx, dy),
            is_left_hand=is_left_hand
        )
        
    def _detect_fingers_fallback(self, image: np.ndarray) -> List[FingerDetection]:
        """
        当YOLO无法使用时的备用手指检测方法
        """
        # 简单的皮肤颜色检测和轮廓分析
        h, w = image.shape[:2]
        
        # 转到HSV色彩空间
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        
        # 皮肤颜色范围 (基于HSV)
        lower_skin = np.array([0, 20, 70], dtype=np.uint8)
        upper_skin = np.array([20, 255, 255], dtype=np.uint8)
        
        # 提取皮肤区域
        skin_mask = cv2.inRange(hsv, lower_skin, upper_skin)
        
        # 形态学操作
        kernel = np.ones((5, 5), np.uint8)
        skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_CLOSE, kernel)
        skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_OPEN, kernel)
        
        # 查找轮廓
        contours, _ = cv2.findContours(skin_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 找到最大的轮廓，假设它是手
        if not contours:
            return []
            
        max_contour = max(contours, key=cv2.contourArea)
        area = cv2.contourArea(max_contour)
        
        # 如果区域太小，不太可能是手
        if area < 5000:
            return []
        
        # 计算轮廓的凸包
        hull = cv2.convexHull(max_contour)
        
        # 使用重心和最远点估算手指位置和方向
        M = cv2.moments(max_contour)
        if M["m00"] == 0:
            return []
            
        cx = int(M["m10"] / M["m00"])
        cy = int(M["m01"] / M["m00"])
        
        # 找到离重心最远的点作为可能的指尖
        max_dist = 0
        tip_x, tip_y = cx, cy
        
        for point in hull:
            x, y = point[0]
            dist = math.sqrt((x - cx)**2 + (y - cy)**2)
            if dist > max_dist:
                max_dist = dist
                tip_x, tip_y = x, y
        
        # 计算方向
        if max_dist > 0:
            dx = (tip_x - cx) / max_dist
            dy = (tip_y - cy) / max_dist
        else:
            dx, dy = 0, 0
        
        # 创建一个模拟的食指检测
        finger_detection = FingerDetection(
            landmarks=[(cx/w, cy/h), (tip_x/w, tip_y/h)],
            finger_type=1,  # 假设为食指
            confidence=0.6,
            tip_position=(tip_x, tip_y),
            direction=(dx, dy),
            is_left_hand=False  # 无法确定，默认为右手
        )
        
        return [finger_detection]

    def get_pointing_direction(
        self,
        finger_detections: List[FingerDetection], 
        prefer_finger_types: List[int] = [1, 2]  # 默认优先选择食指和中指
    ) -> Optional[Tuple[Tuple[int, int], Tuple[float, float]]]:
        """
        获取指向方向，支持左右手
        
        参数:
            finger_detections: 检测到的手指列表
            prefer_finger_types: 优先选择的手指类型
            
        返回:
            Optional[Tuple[Tuple[int, int], Tuple[float, float]]]: (指尖坐标, 方向向量)，如果没有检测到则返回None
        """
        if not finger_detections:
            return None
            
        # 筛选优先手指
        preferred_fingers = [f for f in finger_detections if f.finger_type in prefer_finger_types]
        
        if not preferred_fingers:
            # 如果没有找到优先手指，则使用第一个检测到的手指
            target_finger = finger_detections[0]
        else:
            # 按置信度排序
            preferred_fingers.sort(key=lambda f: f.confidence, reverse=True)
            
            # 如果有多个手指，优先选择指向更明确的手指
            if len(preferred_fingers) > 1:
                # 计算各手指方向向量的大小，选择最大的
                max_magnitude = -1
                target_finger = preferred_fingers[0]
                
                for finger in preferred_fingers:
                    dx, dy = finger.direction
                    magnitude = dx**2 + dy**2
                    if magnitude > max_magnitude:
                        max_magnitude = magnitude
                        target_finger = finger
            else:
                target_finger = preferred_fingers[0]
            
        return (target_finger.tip_position, target_finger.direction)

    def draw_fingers(
        self,
        image: np.ndarray,
        finger_detections: List[FingerDetection],
        draw_landmarks: bool = True,
        draw_direction: bool = True,
        direction_length: int = 100,
        color_right: Tuple[int, int, int] = (0, 0, 255),
        color_left: Tuple[int, int, int] = (255, 0, 0),
        thickness: int = 2
    ) -> np.ndarray:
        """
        在图像上绘制检测到的手指
        
        参数:
            image: 输入图像
            finger_detections: 检测到的手指列表
            draw_landmarks: 是否绘制关键点
            draw_direction: 是否绘制方向线
            direction_length: 方向线长度
            color_right: 右手颜色
            color_left: 左手颜色
            thickness: 线条粗细
            
        返回:
            np.ndarray: 绘制了手指的图像
        """
        result_image = image.copy()
        h, w = result_image.shape[:2]
        
        # 绘制每个手指
        for finger in finger_detections:
            # 选择颜色
            color = color_left if finger.is_left_hand else color_right
            
            # 绘制关键点
            if draw_landmarks and finger.landmarks:
                for i, (x, y) in enumerate(finger.landmarks):
                    # 转换为像素坐标
                    px, py = int(x * w), int(y * h)
                    cv2.circle(result_image, (px, py), 5, color, -1)
                    
                    # 连接点
                    if i > 0:
                        prev_x, prev_y = finger.landmarks[i-1]
                        prev_px, prev_py = int(prev_x * w), int(prev_y * h)
                        cv2.line(result_image, (prev_px, prev_py), (px, py), color, thickness)
            
            # 绘制指尖
            cv2.circle(
                result_image,
                finger.tip_position,
                5,
                color,
                -1  # 填充圆
            )
            
            # 绘制方向线
            if draw_direction:
                end_x = int(finger.tip_position[0] + finger.direction[0] * direction_length)
                end_y = int(finger.tip_position[1] + finger.direction[1] * direction_length)
                cv2.line(
                    result_image,
                    finger.tip_position,
                    (end_x, end_y),
                    color,
                    thickness
                )
                
            # 标注手指类型
            finger_type_names = ['拇指', '食指', '中指', '无名指', '小指']
            finger_name = finger_type_names[finger.finger_type]
            hand_label = "左" if finger.is_left_hand else "右"
            label = f"{hand_label}{finger_name}"
            
            cv2.putText(
                result_image,
                label,
                (finger.tip_position[0] + 10, finger.tip_position[1] - 10),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                color,
                thickness
            )
            
        return result_image